import socket
import struct
import numpy as np
import cv2
import threading
from config import camera_fx, camera_fy, camera_cx, camera_cy

FRAME_WIDTH = 640
FRAME_HEIGHT = 352

class DepthReceiver:
    def __init__(self, host='192.168.1.68', port=5000):
        """ 初始化 TCP 连接 """
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.client_socket.connect((host, port))
        print(f"Connected to {host}:{port}")

        # 线程锁，保证数据同步
        self.lock = threading.Lock()

        # 存储最新的 RGB 和 深度数据
        self.latest_frames = {"bgr": None, "depth": None}

    def recv_all(self, size):
        """ 确保完全接收指定大小的数据 """
        data = b""
        while len(data) < size:
            packet = self.client_socket.recv(size - len(data))
            if not packet:
                return None
            data += packet
        return data

    def receive_frame(self):
        """ 接收 NV12 & 深度数据，确保数据同步 """
        while True:
            # 接收头部信息（数据类型 + 数据大小）
            header = self.recv_all(20)  # 16字节数据类型 + 4字节数据长度
            if not header:
                return None, None  # 连接断开

            data_type, data_size = struct.unpack("16sI", header)
            data_type = data_type.decode().strip("\x00")

            # 接收数据
            data_bytes = self.recv_all(data_size)
            if not data_bytes:
                return None, None

            with self.lock:
                if data_type == "nv12":
                    # 解析 NV12 格式
                    y_size = FRAME_WIDTH * FRAME_HEIGHT
                    y_channel = np.frombuffer(data_bytes[:y_size], dtype=np.uint8).reshape((FRAME_HEIGHT, FRAME_WIDTH))

                    uv_channel = np.frombuffer(data_bytes[y_size:], dtype=np.uint8).reshape((FRAME_HEIGHT // 2, FRAME_WIDTH // 2, 2))
                    u_channel = cv2.resize(uv_channel[:, :, 0], (FRAME_WIDTH, FRAME_HEIGHT), interpolation=cv2.INTER_LINEAR)
                    v_channel = cv2.resize(uv_channel[:, :, 1], (FRAME_WIDTH, FRAME_HEIGHT), interpolation=cv2.INTER_LINEAR)

                    bgr_image = cv2.cvtColor(cv2.merge([y_channel, u_channel, v_channel]), cv2.COLOR_YUV2BGR)

                    # 存储最新的 BGR
                    self.latest_frames["bgr"] = bgr_image

                elif data_type == "depth":
                    # 解析 uint16 深度数据
                    depth_image = np.frombuffer(data_bytes, dtype=np.uint16).reshape((FRAME_HEIGHT, FRAME_WIDTH))

                    # 存储最新的深度图
                    self.latest_frames["depth"] = depth_image

                # 只有当 BGR 和 深度数据都到达时，才返回
                if self.latest_frames["bgr"] is not None and self.latest_frames["depth"] is not None:
                    bgr_frame = self.latest_frames["bgr"]
                    depth_frame = self.latest_frames["depth"]

                    # 清空存储，等待下一个匹配的帧
                    self.latest_frames["bgr"] = None
                    self.latest_frames["depth"] = None
                    return bgr_frame, depth_frame
                else:
                    return None, None

    def display_depth(self, depth_image):
        """ 显示深度图 """
        if depth_image is None:
            return None

        # 归一化并应用伪彩色
        depth_normalized = cv2.normalize(depth_image, None, 0, 255, cv2.NORM_MINMAX)
        depth_8bit = np.uint8(depth_normalized)
        depth_colormap = cv2.applyColorMap(depth_8bit, cv2.COLORMAP_JET)
        return depth_colormap

    def close(self):
        """ 关闭 TCP 连接 """
        self.client_socket.close()
        print("Connection closed.")